import datetime
from enum import Enum
from typing import Union

import pytz


# NOTE we need to make sure all timestamps in UTC
# since most of the providers default to use UTC time
# currently our system is simple but still require us to set the time zone
# either Shanghai or San Francisco, we should discuss this later


class KNOWN_TIMEZONES(str, Enum):
    """Known timezones supported by skywalkers"""

    UTC = "UTC"
    America_San_Francisco = "America/Los_Angeles"
    Asia_Shanghai = "Asia/Shanghai"


def naive_to_aware(
    dt_naive: datetime.datetime,
    tz: Union[KNOWN_TIMEZONES, pytz.BaseTzInfo] = KNOWN_TIMEZONES.America_San_Francisco,
) -> datetime.datetime:
    """Convert a naive `datetime.datetime` object to be time zone aware

    Parameters
    ----------
    dt_naive : datetime.datetime
        A `datetime.datetime` object that is naive, with `dt_naive.tzinfo is None`
    tz : Union[KNOWN_TIMEZONES, pytz.timezone], optional
        Either an item of KNOWN_TIMEZONES, or a `pytz.BaseTzInfo` object,
        by default KNOWN_TIMEZONES.America_San_Francisco

    Returns
    -------
    datetime.datetime
        A `datetime.datetime` object with the same value as the `dt_naive`
        but aware its time zone as specified by the given `tz`.

    Raises
    ------
    ValueError
        When the given `dt_naive` is not a naive `datetime.datetime` object
    """
    if dt_naive.tzinfo is not None:
        raise ValueError(
            "`dt` must be a naive `datetime.datetime`, without a valid tzinfo"
        )
    if isinstance(tz, KNOWN_TIMEZONES):
        time_zone = pytz.timezone(tz.value)
    else:
        time_zone = tz
    return time_zone.localize(dt_naive)


def to_utc_datetime(dt: datetime.datetime) -> datetime.datetime:
    """Convert a time zone aware `datetime.datetime` object to be under the UTC time zone

    Parameters
    ----------
    dt : datetime.datetime
        A `datetime.datetime` object to be converted, with time zone info.

    Returns
    -------
    datetime.datetime
        Converted `datetime.datetime` under the UTC time zone.

    Raises
    ------
    ValueError
        Raised when the given `dt` does not have information about its time zone.
    """
    if dt.tzinfo is None:
        raise ValueError(
            "`dt` must be a timezone aware `datetime.datetime`, with a valid tzinfo"
        )
    else:
        return dt.astimezone(pytz.timezone(KNOWN_TIMEZONES.UTC.value))


def utcnow() -> datetime.datetime:
    """Get the current datetime under the UTC time zone

    Returns
    -------
    datetime.datetime
        current `datetime.datetime` under the UTC time zone.

    Notes
    -----
    It is known that python stdlib's `datetime.datetime.utcnow` is not quite correct,
    and frequently gives out wrong result if used without caution:
    <https://blog.ganssle.io/articles/2019/11/utcnow.html>.
    This function provides a more reliable implementation
    before we can get a better `utcnow` from python stdlib.
    """
    now = datetime.datetime.now()
    return now.astimezone(pytz.timezone(KNOWN_TIMEZONES.UTC.value))
